#include "common.h"
#include "kernel_operator.h"
#ifdef __CCE_KT_TEST__
#include <iostream>
#endif
namespace AscendC
{
template <typename _DT_PREDICT, typename _DT_LABEL, typename _DT_Y> class KernelMean
{
  public:
    TPipe pipe;
    CommonTiling tilingData;

    DefInTensor(PREDICT);
    DefInTensor(LABEL);
    DefGlobalTensor(Y);
    DefBufVECOUT(TEMP);

  public:
    __aicore__ inline KernelMean()
    {
    }
    __aicore__ inline void Init(GM_ADDR predict, GM_ADDR label, GM_ADDR y, uint64_t size, CommonTiling &tiling)
    {
        tilingData = tiling;

        // 切分global
        SetGlobalBuffer(PREDICT, predict, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(LABEL, label, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, 0, 32 / sizeof(_DT_PREDICT));
        // 初始化队列
        InitQueueSimple(PREDICT, tilingData.bufferSize);
        InitQueueSimple(LABEL, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(TEMP, 32);
        // 将TEMP置为0
        TBufGet(TEMP, _DT_PREDICT);
        Muls(TBufTensorName(TEMP), TBufTensorName(TEMP), _DT_PREDICT(0), 1);
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < finnal_progress; ++i)
        {
            CopyIn(i * tilingData.bufferSize, tilingData.bufferSize);
            Compute(tilingData.bufferSize, tilingData.bufferSize);
        }
        auto copy_count = ALIGN_TO(tilingData.last_size, 32 / sizeof(_DT_PREDICT));
        CopyIn(finnal_progress * tilingData.bufferSize, copy_count);
        ComputeFinal(tilingData.finalCalcCount);
    }
    __aicore__ inline void CopyIn(uint32_t index, uint32_t calcCount)
    {
        EnQueGlobal2Local(PREDICT, index, calcCount);
        EnQueGlobal2Local(LABEL, index, calcCount);
    }

    __aicore__ inline void Compute(uint32_t calcCount, uint32_t sum_count)
    {
        DeQueSimple(PREDICT);
        DeQueSimple(LABEL);
        TBufGet(TEMP, _DT_PREDICT);
        // predict - label
        Sub(LTensorName(PREDICT), LTensorName(PREDICT), LTensorName(LABEL), calcCount);
        // (predict - label) * (predict - label)
        Mul(LTensorName(LABEL), LTensorName(PREDICT), LTensorName(PREDICT), calcCount);
        // 求和
        ReduceSum(LTensorName(PREDICT), LTensorName(LABEL), LTensorName(LABEL), sum_count);
        // 加到总数上
        Add(TBufTensorName(TEMP), TBufTensorName(TEMP), LTensorName(PREDICT), 1);
        QueFree(LABEL);
        QueFree(PREDICT);
    }

    __aicore__ inline void ComputeFinal(uint32_t calcCount)
    {
        Compute(calcCount, tilingData.last_size);
        TBufGet(TEMP, _DT_PREDICT);
        Muls(TBufTensorName(TEMP), TBufTensorName(TEMP), _DT_PREDICT(1 / float(tilingData.size)), 1);
        DataCopy(GTensorName(Y), TBufTensorName(TEMP), 32 / sizeof(TypeOf(PREDICT)));
    }
};

template <typename _DT_PREDICT, typename _DT_LABEL, typename _DT_Y> class KernelSum
{
  public:
    TPipe pipe;
    CommonTiling tilingData;

    DefInTensor(PREDICT);
    DefInTensor(LABEL);
    DefOutTensor(Y);

    DefBufVECIN(WORK_LOCAL);
    DefBufVECIN(TEMP);

  public:
    __aicore__ inline KernelSum()
    {
    }
    __aicore__ inline void Init(GM_ADDR predict, GM_ADDR label, GM_ADDR y, uint64_t size, CommonTiling &tiling)
    {
        tilingData = tiling;

        // 切分global
        SetGlobalBuffer(PREDICT, predict, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(LABEL, label, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(PREDICT, tilingData.bufferSize);
        InitQueueSimple(LABEL, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
        // 初始化buf
        InitTBufBuffer(WORK_LOCAL, tilingData.bufferSize * sizeof(_DT_PREDICT));
        InitTBufBuffer(TEMP, tilingData.bufferSize * sizeof(_DT_PREDICT));
        // 将TEMP置为0
        TBufGet(TEMP, _DT_PREDICT);
        for (auto i = 0; i < tilingData.bufferSize; ++i)
        {
            TBufTensorName(TEMP).SetValue(i, _DT_PREDICT(0));
        }
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < finnal_progress; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(tilingData.bufferSize, false);
        }
        auto copy_count = tilingData.getFinalCopyCount<TypeOf(PREDICT)>();
        CopyIn(finnal_progress, copy_count);
        Compute(tilingData.finalCalcCount, true);
        CopyOut();
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(PREDICT, index, calcCount);
        EnQueGlobal2Local(LABEL, index, calcCount);
    }
    __aicore__ inline void Compute(uint32_t calcCount, bool final)
    {
        DeQueSimple(PREDICT);
        DeQueSimple(LABEL);
        TBufGet(WORK_LOCAL, _DT_PREDICT);
        TBufGet(TEMP, _DT_PREDICT);
        // predict - label
        Sub(LTensorName(PREDICT), LTensorName(PREDICT), LTensorName(LABEL), calcCount);
        // (predict - label) * (predict - label)
        Mul(LTensorName(LABEL), LTensorName(PREDICT), LTensorName(PREDICT), calcCount);
        // // 计算均值
        int32_t count = final ? tilingData.last_size : calcCount;
        ReduceSum(LTensorName(PREDICT), LTensorName(LABEL), TBufTensorName(WORK_LOCAL), count);
        // 加到总数上
        if (final)
        {
            QueAllocSimple(Y);
            Add(LTensorName(Y), TBufTensorName(TEMP), LTensorName(PREDICT), calcCount);
            EnQue(Y);
        }
        else
        {
            Add(TBufTensorName(TEMP), TBufTensorName(TEMP), LTensorName(PREDICT), calcCount);
        }
        QueFree(LABEL);
        QueFree(PREDICT);
    }
    __aicore__ inline void CopyOut()
    {
        DeQueLocal2Global(Y, 0, ALIGN_TO(1, 32 / sizeof(TypeOf(PREDICT))));
    }
};

template <typename _DT_PREDICT, typename _DT_LABEL, typename _DT_Y> class KernelNone
{
  public:
    TPipe pipe;
    CommonTiling tilingData;

    DefInTensor(PREDICT);
    DefInTensor(LABEL);
    DefOutTensor(Y);

  public:
    __aicore__ inline KernelNone()
    {
    }
    __aicore__ inline void Init(GM_ADDR predict, GM_ADDR label, GM_ADDR y, CommonTiling &tiling)
    {
        tilingData = tiling;
        // 切分global
        SetGlobalBuffer(PREDICT, predict, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(LABEL, label, tilingData.startIdx, tilingData.blockLength);
        SetGlobalBuffer(Y, y, tilingData.startIdx, tilingData.blockLength);
        // 初始化队列
        InitQueueSimple(PREDICT, tilingData.bufferSize);
        InitQueueSimple(LABEL, tilingData.bufferSize);
        InitQueueSimple(Y, tilingData.bufferSize);
    }
    __aicore__ inline void Process()
    {
        auto finnal_progress = tilingData.loopCount - 1;
        for (uint32_t i = 0; i < finnal_progress; ++i)
        {
            CopyIn(i, tilingData.bufferSize);
            Compute(i, tilingData.bufferSize);
            CopyOut(i, tilingData.bufferSize);
        }
        auto copy_count = tilingData.getFinalCopyCount<TypeOf(PREDICT)>();
        CopyIn(finnal_progress, copy_count);
        Compute(finnal_progress, tilingData.finalCalcCount);
        CopyOut(finnal_progress, copy_count);
    }
    __aicore__ inline void CopyIn(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        EnQueGlobal2Local(PREDICT, index, calcCount);
        EnQueGlobal2Local(LABEL, index, calcCount);
    }
    __aicore__ inline void Compute(uint32_t progress, uint32_t calcCount)
    {
        DeQueSimple(PREDICT);
        DeQueSimple(LABEL);
        QueAllocSimple(Y);
        // predict - label
        Sub(LTensorName(PREDICT), LTensorName(PREDICT), LTensorName(LABEL), calcCount);
        // (predict - label) * (predict - label)
        Mul(LTensorName(Y), LTensorName(PREDICT), LTensorName(PREDICT), calcCount);
        EnQue(Y);
        QueFree(LABEL);
        QueFree(PREDICT);
    }
    __aicore__ inline void CopyOut(uint32_t progress, uint32_t calcCount)
    {
        auto index = progress * tilingData.bufferSize;
        DeQueLocal2Global(Y, index, calcCount);
    }
};
} // namespace AscendC
extern "C" __global__ __aicore__ void mse_loss(GM_ADDR predict, GM_ADDR label, GM_ADDR y, GM_ADDR workspace,
                                               GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    CRATE_COMMON_TILING_SIMPLE;
    switch (tiling_data.tag)
    {
    case 0: {
        AscendC::KernelMean<DTYPE_PREDICT, DTYPE_LABEL, DTYPE_Y> op;
        op.Init(predict, label, y, tiling_data.size, commonTiling);
        op.Process();
    }
    break;
    case 1: {
        AscendC::KernelSum<DTYPE_PREDICT, DTYPE_LABEL, DTYPE_Y> op;
        op.Init(predict, label, y, tiling_data.size, commonTiling);
        op.Process();
    }
    break;
    case 2: {
        AscendC::KernelNone<DTYPE_PREDICT, DTYPE_LABEL, DTYPE_Y> op;
        op.Init(predict, label, y, commonTiling);
        op.Process();
    }
    break;
    }
}